package geppetto.cat.programs;


import geppetto.cat.alignments.Alignment;
import geppetto.cat.alignments.AlignmentEvaluator;
import geppetto.cat.alignments.AlignmentSymetrization;
import geppetto.cat.alignments.AlignmentsSet;
import geppetto.cat.common.StaticTools;
import geppetto.cat.corpus.BilingualCorpus;
import geppetto.cat.models.AbstractModel;
import geppetto.cat.models.AgreementHMM;
import geppetto.cat.models.HMM;
import geppetto.cat.models.SubstochasticHMM;

import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.PrintStream;
import java.util.Random;



public class AlignmentsForMosesPosteriorBased {
	String corpusDescription;

	int size, maxSentenceLen;
	String forwardModelDir;
	String backwardModelDir;
	String modelName;
	String baseDir;
	boolean streamInitialized = false;

	double[] posts = {0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9};
	PrintStream[][] forwardStream, backwardStream;
	BilingualCorpus forwardCorpus, backwardCorpus;



	
	double epsilon; 
	double slack;
	int maxSteps;
	int maxProjectionSteps;
	boolean projectDecoding;
	
	public AlignmentsForMosesPosteriorBased(String[] args) {
		corpusDescription = args[0];
		size = Integer.parseInt(args[1]);
		maxSentenceLen = Integer.parseInt(args[2]);
		forwardModelDir = args[3];
		backwardModelDir = args[4];
		modelName = args[5];
		baseDir = args[6];
		 
		System.out.println("Corpus " + corpusDescription);
		System.out.println("Size " + size);
		System.out.println("Max Sentence Len " + maxSentenceLen);
		System.out.println("Forward odelDir " + forwardModelDir);
		System.out.println("Backward odelDir " + backwardModelDir);
		System.out.println("Model Name " + modelName);
		System.out.println("Output Dir " + baseDir);
		System.out.println("---");
		
		if (modelName.equalsIgnoreCase("symmetric") || modelName.equalsIgnoreCase("bijective")){
			epsilon = Double.parseDouble(args[7]);
			slack = Double.parseDouble(args[8]);
			maxSteps = Integer.parseInt(args[9]);
			maxProjectionSteps = Integer.parseInt(args[10]);
			projectDecoding = Boolean.parseBoolean(args[11]);
			System.out.println("epsilon " + epsilon);
			System.out.println("slack " + slack);
			System.out.println("maxSteps " + maxSteps);
			System.out.println("maxProjectIterations " + maxProjectionSteps);
			System.out.println("projectDecoding " + projectDecoding);
		}		
		forwardStream = new PrintStream[posts.length][3];
		backwardStream = new PrintStream[posts.length][3];
	}

	public void initializeCorpus(String corpusDesc, int size)
			throws IOException {
		System.out.println("Loading Corpus");
		forwardCorpus = BilingualCorpus.getCorpusFromFileDescription(
				corpusDesc, size, maxSentenceLen);
		backwardCorpus = forwardCorpus.reverse();
		System.out.println("Finished Loading Corpus");
	}

	

	void initializeStreams(String modelName, String baseDir,
			String decodingType, String sourcePrefix, String foreignPrefix,
			String corpusName) throws FileNotFoundException, IOException {
		System.out.println("Initializing Streams");
		initializeStream(forwardStream, backwardStream, sourcePrefix,
				foreignPrefix, "softUnion", modelName, baseDir, decodingType,
				corpusName);
	}

	void initializeStream(PrintStream[][] forwardStream,
			PrintStream[][] backwardStream, String sourcePrefix,
			String foreignPrefix, String alignmentName, String modelName,
			String baseDir, String decodingType, String corpusName)
			throws IOException, FileNotFoundException {
		
		
		for(int postIndex =0; postIndex < posts.length; postIndex++){
			String forwardDirection = baseDir + "/" + modelName + "/" + size + "/"
			+ decodingType + "-" + projectDecoding+ "-" + posts[postIndex] + "/" 
			+ alignmentName + "/" + corpusName + "/"
			+ sourcePrefix + "-" + foreignPrefix + "/model/";
	String backwardDirection = baseDir + "/" + modelName + "/"  + size + "/"
			+ decodingType + "-" + projectDecoding+ "-" + posts[postIndex] + "/" 
			+ alignmentName + "/" + corpusName + "/"
			+ foreignPrefix + "-" + sourcePrefix + "/model/";
	StaticTools.createDir(forwardDirection);
	StaticTools.createDir(backwardDirection);
			forwardStream[postIndex][0] = new PrintStream(new FileOutputStream(
					forwardDirection + "aligned." + alignmentName), true, "UTF-8");
			forwardStream[postIndex][1] = new PrintStream(new FileOutputStream(
				forwardDirection + "aligned.0." + sourcePrefix), true, "UTF-8");
			forwardStream[postIndex][2] = new PrintStream(new FileOutputStream(
				forwardDirection + "aligned.0." + foreignPrefix), true, "UTF-8");
			backwardStream[postIndex][0] = new PrintStream(new FileOutputStream(
					backwardDirection + "aligned." + alignmentName), true, "UTF-8");
			backwardStream[postIndex][1] = new PrintStream(new FileOutputStream(
				backwardDirection + "aligned.0." + foreignPrefix), true,
			"UTF-8");
			backwardStream[postIndex][2] = new PrintStream(new FileOutputStream(
				backwardDirection + "aligned.0." + sourcePrefix), true, "UTF-8");
		}
	}

	void saveAlignment(PrintStream[] forwardStream,
			PrintStream[] backwardStream, Alignment forward,
			Alignment backward, BilingualCorpus corpus,
			BilingualCorpus backCorpus) throws IOException,
			FileNotFoundException {
		String[] data = forward.toMosesFormat(corpus);
		forwardStream[0].println(data[0]);
		forwardStream[1].println(data[1]);
		forwardStream[2].println(data[2]);
		data = backward.toMosesFormat(backCorpus);
		backwardStream[0].println(data[0]);
		backwardStream[1].println(data[1]);
		backwardStream[2].println(data[2]);
	}

	

	
	
	public void run() throws IOException {
		initializeCorpus(corpusDescription, size);
		AbstractModel models[] = GenericMethods.initializeModelsFullDir(forwardCorpus, backwardCorpus, modelName, forwardModelDir,backwardModelDir, epsilon, slack, maxSteps, maxProjectionSteps);
		AbstractModel forward = models[0];
		AbstractModel backward = models[1];
		initializeStreams(modelName, baseDir, "post-fix", forwardCorpus
				.getSourceSufix(), forwardCorpus.getForeignSufix(),
				forwardCorpus.getName());
		for (int i = 0; i < forwardCorpus.getNumSentences(BilingualCorpus.TRAIN_CORPUS); i++) {
				Alignment f = forward.posteriorDecodingAlignment(i,
						BilingualCorpus.TRAIN_CORPUS, 0, projectDecoding, null);
				Alignment b = backward.posteriorDecodingAlignment(i,
						BilingualCorpus.TRAIN_CORPUS, 0, projectDecoding, null);				
				for(int postIndex = 0; postIndex < posts.length; postIndex++){
					double postTresh = posts[postIndex];
					Alignment softUnionF = AlignmentSymetrization.softUnion(f, b, postTresh);
					Alignment softUnionB = softUnionF.reverse();
					saveAlignment(forwardStream[postIndex], backwardStream[postIndex], softUnionF, softUnionB, forwardCorpus, backwardCorpus);		
				}
				if (i % 1000 == 0){
					System.out.print("....." + i);
				}
		}
		 	
	}

	public void closeStream() {
		System.out.println("Closing Streams");
		for(int postIndex = 0; postIndex < posts.length; postIndex++){
			for (int i = 0; i < 3; i++) {
				forwardStream[postIndex][i].close();
				backwardStream[postIndex][i].close();
			}
		}
	}

	public static void main(String[] args) throws IOException {
		AlignmentsForMosesPosteriorBased exp = new AlignmentsForMosesPosteriorBased(args);
		exp.run();
		exp.closeStream();
	}

}
